"""
投影操作符实现
"""

from typing import Optional, List, Dict, Any, Iterator, Callable
import re
import time

from ..executor_base import IteratorExecutor
from ..types import Record, RecordBatch, ExecutionContext, DataValue

class ProjectOperator(IteratorExecutor):
    """投影操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor, columns: List[str]):
        super().__init__(executor_id, context)
        self.child = child
        self.columns = columns
        self.computed_columns: Dict[str, Callable[[Record], Any]] = {}
        
        # 添加子执行器
        self.add_child(child)
        
        # 解析列规格
        self._parse_columns()
    
    def _parse_columns(self):
        """解析列规格，支持简单的计算列"""
        for column in self.columns:
            if '=' in column:
                # 计算列格式: "alias = expression"
                parts = column.split('=', 1)
                if len(parts) == 2:
                    alias = parts[0].strip()
                    expression = parts[1].strip()
                    self.computed_columns[alias] = self._compile_expression(expression)
    
    def _compile_expression(self, expression: str) -> Callable[[Record], Any]:
        """编译表达式为可执行函数"""
        expression = expression.strip()
        
        # 简单的表达式支持
        if '+' in expression:
            # 加法表达式
            parts = expression.split('+')
            if len(parts) == 2:
                left_expr = parts[0].strip()
                right_expr = parts[1].strip()
                
                def add_func(record: Record) -> Any:
                    left_val = self._evaluate_expression(left_expr, record)
                    right_val = self._evaluate_expression(right_expr, record)
                    try:
                        return float(left_val) + float(right_val)
                    except (ValueError, TypeError):
                        return None
                
                return add_func
        
        elif '-' in expression:
            # 减法表达式
            parts = expression.split('-')
            if len(parts) == 2:
                left_expr = parts[0].strip()
                right_expr = parts[1].strip()
                
                def sub_func(record: Record) -> Any:
                    left_val = self._evaluate_expression(left_expr, record)
                    right_val = self._evaluate_expression(right_expr, record)
                    try:
                        return float(left_val) - float(right_val)
                    except (ValueError, TypeError):
                        return None
                
                return sub_func
        
        elif '*' in expression:
            # 乘法表达式
            parts = expression.split('*')
            if len(parts) == 2:
                left_expr = parts[0].strip()
                right_expr = parts[1].strip()
                
                def mul_func(record: Record) -> Any:
                    left_val = self._evaluate_expression(left_expr, record)
                    right_val = self._evaluate_expression(right_expr, record)
                    try:
                        return float(left_val) * float(right_val)
                    except (ValueError, TypeError):
                        return None
                
                return mul_func
        
        elif '/' in expression:
            # 除法表达式
            parts = expression.split('/')
            if len(parts) == 2:
                left_expr = parts[0].strip()
                right_expr = parts[1].strip()
                
                def div_func(record: Record) -> Any:
                    left_val = self._evaluate_expression(left_expr, record)
                    right_val = self._evaluate_expression(right_expr, record)
                    try:
                        left_num = float(left_val)
                        right_num = float(right_val)
                        if right_num != 0:
                            return left_num / right_num
                        else:
                            return None
                    except (ValueError, TypeError):
                        return None
                
                return div_func
        
        elif expression.upper().startswith('UPPER(') and expression.endswith(')'):
            # UPPER函数
            inner_expr = expression[6:-1].strip()
            
            def upper_func(record: Record) -> Any:
                value = self._evaluate_expression(inner_expr, record)
                if value is not None:
                    return str(value).upper()
                return None
            
            return upper_func
        
        elif expression.upper().startswith('LOWER(') and expression.endswith(')'):
            # LOWER函数
            inner_expr = expression[6:-1].strip()
            
            def lower_func(record: Record) -> Any:
                value = self._evaluate_expression(inner_expr, record)
                if value is not None:
                    return str(value).lower()
                return None
            
            return lower_func
        
        elif expression.upper().startswith('CONCAT(') and expression.endswith(')'):
            # CONCAT函数
            inner_expr = expression[7:-1].strip()
            args = [arg.strip() for arg in inner_expr.split(',')]
            
            def concat_func(record: Record) -> Any:
                result_parts = []
                for arg in args:
                    value = self._evaluate_expression(arg, record)
                    if value is not None:
                        result_parts.append(str(value))
                return ''.join(result_parts)
            
            return concat_func
        
        else:
            # 简单表达式（列名或常量）
            def simple_func(record: Record) -> Any:
                return self._evaluate_expression(expression, record)
            
            return simple_func
    
    def _evaluate_expression(self, expr: str, record: Record) -> Any:
        """计算表达式值"""
        expr = expr.strip()
        
        # 如果是列名
        if expr in record:
            return record[expr]
        
        # 处理带表前缀的列名 (如 left.name, right.id)
        if '.' in expr:
            return record.get(expr)
        
        # 如果是字符串字面量
        if expr.startswith("'") and expr.endswith("'"):
            return expr[1:-1]
        
        if expr.startswith('"') and expr.endswith('"'):
            return expr[1:-1]
        
        # 如果是数字字面量
        try:
            if '.' in expr:
                return float(expr)
            else:
                return int(expr)
        except ValueError:
            pass
        
        # 布尔字面量
        if expr.lower() == 'true':
            return True
        elif expr.lower() == 'false':
            return False
        elif expr.lower() == 'null':
            return None
        
        # 如果都不是，返回原字符串
        return expr
    
    def open(self):
        """打开投影操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.child.open()
            
            # 创建投影迭代器
            self.current_iterator = self._create_project_iterator()
    
    def _create_project_iterator(self) -> Iterator[Record]:
        """创建投影迭代器"""
        while True:
            child_batch = self.child.next_batch()
            if not child_batch:
                break
            
            projected_batch = []
            
            for record in child_batch:
                self.check_timeout()
                
                # 更新处理统计
                self.update_stats(rows_processed=1)
                
                # 执行投影
                projected_record = self._project_record(record)
                projected_batch.append(projected_record)
                
                # 更新返回统计
                self.update_stats(rows_returned=1)
            
            # 返回投影后的记录
            for record in projected_batch:
                yield record
    
    def _project_record(self, record: Record) -> Record:
        """对单条记录执行投影"""
        projected = {}
        
        for column in self.columns:
            if '=' in column:
                # 计算列
                alias = column.split('=', 1)[0].strip()
                if alias in self.computed_columns:
                    try:
                        projected[alias] = self.computed_columns[alias](record)
                    except Exception:
                        projected[alias] = None
            else:
                # 普通列
                column = column.strip()
                
                # 处理 * 通配符
                if column == '*':
                    # 添加所有列
                    projected.update(record)
                elif '.' in column:
                    # 带前缀的列名
                    projected[column] = record.get(column)
                else:
                    # 简单列名
                    if column in record:
                        projected[column] = record[column]
                    else:
                        # 尝试查找带前缀的列
                        found = False
                        for key in record.keys():
                            if key.endswith('.' + column):
                                projected[column] = record[key]
                                found = True
                                break
                        
                        if not found:
                            projected[column] = None
        
        return projected
    
    def close(self):
        """关闭投影操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.child.close()
            
            self.is_open = False
            self.current_iterator = None
    
    def get_output_columns(self) -> List[str]:
        """获取输出列列表"""
        output_columns = []
        
        for column in self.columns:
            if '=' in column:
                # 计算列的别名
                alias = column.split('=', 1)[0].strip()
                output_columns.append(alias)
            elif column.strip() == '*':
                # 通配符需要根据输入确定
                output_columns.append('*')
            else:
                output_columns.append(column.strip())
        
        return output_columns

class ComputedColumnOperator(ProjectOperator):
    """计算列操作符，专门处理复杂的计算列"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor, 
                 computed_columns: Dict[str, str]):
        # 将计算列转换为投影格式
        columns = []
        for alias, expression in computed_columns.items():
            columns.append(f"{alias} = {expression}")
        
        super().__init__(executor_id, context, child, columns)

class RenameOperator(ProjectOperator):
    """重命名操作符，用于列重命名"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor,
                 column_mapping: Dict[str, str]):
        # 将重命名映射转换为投影格式
        columns = []
        for old_name, new_name in column_mapping.items():
            columns.append(f"{new_name} = {old_name}")
        
        super().__init__(executor_id, context, child, columns)

